from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import torch
import socket
import numpy as np
import cv2
import time
from threading import Thread  # 多线程实现Socket通信
# 高并发情况下，粘包真是烦死了烦死了QAQ
# Socket就用来激活一下程序来读取文件，虽然这样做烧点磁盘IO..
train_new_model = False;
should_i_test = False;
host = 'localhost';
port = 14159;
training_epoch = 13;
client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
datas = []
class MNISTModel(torch.nn.Module):
    def __init__(self):
        super(MNISTModel, self).__init__()
        self.linear1 = torch.nn.Linear(784,512);
        self.linear2 = torch.nn.Linear(512, 256);
        self.linear3 = torch.nn.Linear(256, 128);
        self.linear4 = torch.nn.Linear(128, 64);
        self.linear5 = torch.nn.Linear(64, 32);
        self.linear6 = torch.nn.Linear(32, 10);
        self.activate = torch.nn.ReLU();

    def forward(self,x):
        x = x.view(-1,784);
        x = self.activate(self.linear1(x));
        x = self.activate(self.linear2(x));
        x = self.activate(self.linear3(x));
        x = self.activate(self.linear4(x));
        x = self.activate(self.linear5(x));
        x = (self.linear6(x));
        return x;

def connect():
    client.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)  # 在客户端开启心跳维护
    try:
        client.connect((host, port))
        print("模型已成功连接至服务器")
    except Exception:
        print("连接失败，三秒后重连")
        time.sleep(3)


def receive(model,datas):   #接收数据并返回
        data = client.recv(1024);

        if "OK" not in data.decode():
            for i in data.decode():
                datas.append(np.float32(i));
        else:
            print("数据接收完成")
            x = torch.from_numpy(np.asarray(datas)).view(1,1,784,512)
            y_hat = model(x);
            _, predited = torch.max(y_hat.data, dim=1);
            for index, y_pred in enumerate(predited.detach().numpy().tolist(), 0):
                print(y_pred);
            #print(predited.detach().numpy().tolist())
                datas.clear();
                client.send(str((predited.detach().numpy())).encode("UTF-8"));

#  数据集的初始化工作
batch_size=64
transform = transforms.Compose([  # 归一化的东西 Compose 把中括号的对象转为张量
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081, ))  # 正态分布的归一化
])

train_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=True,
                               download=True,
                               transform=transform);

train_loader = DataLoader(train_dataset,shuffle=True,batch_size=batch_size);
test_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=False,
                               download=True,
                               transform=transform);
test_loader = DataLoader(test_dataset,shuffle=True,batch_size=batch_size);
# 数据集的初始化工作

# 模型/训练


if train_new_model:
    model = MNISTModel();
    criteria = torch.nn.CrossEntropyLoss();  # 交叉熵
    optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.15);  # momentum 是冲量，作用：更快冲出局部最低点

    for epoch in range(13):
        for batch_id,data in enumerate(train_loader,0):  # 供给数据
            x,y=data  # 解构参数
            y_hat = model(x);
            loss = criteria(y_hat,y);
            optimizer.zero_grad();
            loss.backward();
            optimizer.step();
            print("epoch:", epoch,"batch:",batch_id, "loss:", loss.item());
    # 模型/训练
    torch.save(model,"HandWritingNumberModel.pkl")  # 持久化
else:
    print("加载模型...")
    model = torch.load("HandWritingNumberModel.pkl");
    print("加载模型完成!")
    # 测试
    if should_i_test:
        total=0;
        correct=0;
        for data in test_loader:
            x,y = data;
            y_hat = model(x);
            _, predited = torch.max(y_hat.data, dim=1);
            #print(predited.detach().numpy().tolist())
            total+=y.size(0)
            for index,y_pred in enumerate(predited.detach().numpy().tolist(),0):
                y_true = y.detach().numpy().tolist();
                print("真实值",y_true[index],"预测值",y_pred);
                if y_true[index] == y_pred:
                    correct+=1;
        print("正确率:",100*correct/total);
        # 测试
    connect();
    while True:
        try:
            rec_thread = Thread(target=receive(model=model,datas=datas));
            rec_thread.start();
        except socket.error:
            print("模型与服务器断开连接，正在重新连接")
            connect();  # 自动重连
        #except Exception as e:
        #    print("未知错误")
        #    connect();

            # threading.Thread(target=receive);
